#include "xbyak/xbyak.h"
#include "blk.h"
#include "mach.h"
#include "helper.h"
#include "exec.h"

extern "C" void * jit_exec_malloc (size_t size);
extern "C" void * jit_translate (void * blk, void *exec_ptr, size_t size);

class LC3Generator : public Xbyak::CodeGenerator {
    void operator=(const LC3Generator&);
    public:
        LC3Generator (BasicBlock *blk, void *userPtr = 0, size_t size = Xbyak::DEFAULT_MAX_CODE_SIZE) :
        Xbyak::CodeGenerator(size, userPtr) {
            MachineStatus *ms;
            // rdi <- ms->register
            // rsi <- ms->memory
            // rdx <- helper selector
            /*
             * R0 rdi + 0
             * R1 rdi + 2
             * R2 rdi + 4
             * R3 rdi + 6
             * R4 rdi + 8
             * R5 rdi + 10
             * R6 rdi + 12
             * R7 rdi + 14
             * IR rdi + 16
             * PC rdi + 18
             * PSR rdi + 20
             * 
             * PC always in rdx
             * IR always in rcx
             * r8, r9, r10 r11 are scratch registers
             */
            ms = (MachineStatus *) blk->machine_status;
            target_size_t *insn_ptr;
            target_size_t *registers = ms->registers;
            #define INST (*insn_ptr)
            mov (rdx, word [rdi + 18]);
            and_ (rdx, 0xffff);
            for (insn_ptr = (target_size_t*) blk->lc3_insn_blk.start; 
                    insn_ptr < (target_size_t*) blk->lc3_insn_blk.end; insn_ptr++) {
                inLocalLabel ();
                // load mem[PC] to IR, note that lc3 addressing by word (2 bytes)
                mov (rcx, word [rsi + rdx * 2]);
                and_ (rcx, 0xffff);
                // now update pc
                add (rdx, 1);
                // IR and PC updated
                DEF_INST (ADD, *insn_ptr, 0xF038, 0x1000, {
                    mov (r8w, word [ rdi + (SR1 * 2)]);
                    mov (r9w, word [ rdi + (SR2 * 2)]);
                    add (r8w, r9w);
                    mov (word [rdi + (DR * 2)], r8w);
                    xor_ (r9w, r9w);
                    test (r8w, 0x8000);
                    jnz (".n");
                    cmp (r8w, r9w);
                    jz (".z");
                    jnz (".p");
                    L(".n");
                    add (r9w, 4);
                    jmp (".out");
                    L(".z");
                    add (r9w, 2);
                    jmp (".out");
                    L(".p");
                    add (r9w, 1);
                    jmp (".out");
                    L(".out");
                    mov (word [rdi + 20], r9w);
                }); 

                DEF_INST (ADDI,*insn_ptr, 0xF020, 0x1020, {
                    mov (r8w, word [ rdi + (SR1 * 2)]);
                    mov (r9w, imm5);
                    add (r8w, r9w);
                    mov (word [rdi + (DR * 2)], r8w);
                    xor_ (r9w, r9w);
                    test (r8w, 0x8000);
                    jnz (".n");
                    cmp (r8w, r9w);
                    jz (".z");
                    jnz (".p");
                    L(".n");
                    add (r9w, 4);
                    jmp (".out");
                    L(".z");
                    add (r9w, 2);
                    jmp (".out");
                    L(".p");
                    add (r9w, 1);
                    jmp (".out");
                    L(".out");
                    mov (word [rdi + 20], r9w);
                });

                DEF_INST (AND,*insn_ptr, 0xF038, 0x5000, {
                    mov (r8w, word [ rdi + (SR1 * 2)]);
                    mov (r9w, word [ rdi + (SR2 * 2)]);
                    and_ (r8w, r9w);
                    mov (word [rdi + (DR * 2)], r8w);
                    xor_ (r9w, r9w);
                    test (r8w, 0x8000);
                    jnz (".n");
                    cmp (r8w, r9w);
                    jz (".z");
                    jnz (".p");
                    L(".n");
                    add (r9w, 4);
                    jmp (".out");
                    L(".z");
                    add (r9w, 2);
                    jmp (".out");
                    L(".p");
                    add (r9w, 1);
                    jmp (".out");
                    L(".out");
                    mov (word [rdi + 20], r9w);
                });

                DEF_INST (ANDI,*insn_ptr, 0xF020, 0x5020, {
                    mov (r8w, word [ rdi + (SR1 * 2)]);
                    mov (r9w, imm5);
                    and_ (r8w, r9w);
                    mov (word [rdi + (DR * 2)], r8w);
                    xor_ (r9w, r9w);
                    test (r8w, 0x8000);
                    jnz (".n");
                    cmp (r8w, r9w);
                    jz (".z");
                    jnz (".p");
                    L(".n");
                    add (r9w, 4);
                    jmp (".out");
                    L(".z");
                    add (r9w, 2);
                    jmp (".out");
                    L(".p");
                    add (r9w, 1);
                    jmp (".out");
                    L(".out");
                    mov (word [rdi + 20], r9w);
                });

                DEF_INST (BR,*insn_ptr, 0xF000, 0x0000, {
                    mov (r8w, word [rdi + 20]);
                    mov (r9w, FLAGS);
                    test (r8w, r9w);
                    jz (".out");
                    add (dx, imm9);
                    L (".out");
                    mov (word [rdi + 18], dx);
                    mov (rax, rdx);
                    break;
                });

                DEF_INST (JMP,*insn_ptr, 0xFE3F, 0xC000, {
                    mov (dx, word [rdi + BaseR * 2]);
                    // syn pc
                    mov (word [rdi + 18], dx);
                    // retval is next block address
                    mov (rax, rdx);
                    break;
                });

                DEF_INST (JSR,*insn_ptr, 0xF800, 0x4800, {
                    mov (word [rdi + 14], dx);
                    mov (r8w, dx);
                    add (r8w, imm11);
                    and_ (r8w, 0xffff);
                    mov (word [rdi + 18], r8w);
                    mov (rdx, r8);
                    mov (rax, r8);
                    // ret ();
                    break;
                });

                DEF_INST (JSRR,*insn_ptr, 0xFE3F, 0x4000, {
                    mov (r8w, word [rdi + BaseR * 2]);
                    mov (word [rdi + 14], dx);
                    mov (word [rdi + 18], r8w);
                    mov (rdx, r8);
                    mov (rax, r8);
                });

                DEF_INST (LD, *insn_ptr,0xF000, 0x2000, {
                    mov (r8, imm9);
                    add (r8, rdx);
                    mov (r8w, word [rsi +r8 * 2]);
                    mov (word [rdi + DR * 2], r8w);
                    xor_ (r9w, r9w);
                    test (r8w, 0x8000);
                    jnz (".n");
                    cmp (r8w, r9w);
                    jz (".z");
                    jnz (".p");
                    L(".n");
                    add (r9w, 4);
                    jmp (".out");
                    L(".z");
                    add (r9w, 2);
                    jmp (".out");
                    L(".p");
                    add (r9w, 1);
                    jmp (".out");
                    L(".out");
                    mov (word [rdi + 20], r9w);
                });

                DEF_INST (LDI,*insn_ptr, 0xF000, 0xA000, {
                     //read_memory ((REG (R_PC) + imm9) & 0xFFFF, &(REG (DR)));
                    push (rdi);
                    push (rsi);
                    push (rdx);
                    push (rcx);
                    // read_memory : (reg ms pc insn)->memrd
                    mov (rsi, (uint64_t)ms);
                    call (helper_ldi);
                    pop (rcx);
                    pop (rdx);
                    pop (rsi);
                    pop (rdi);
                });

                DEF_INST (LDR,*insn_ptr, 0xF000, 0x6000, {
                    mov (r8, imm6);
                    mov (r9w, word [rdi + BaseR * 2]);
                    add (r8, r9);
                    mov (r8w, word [rsi +r8 * 2]);
                    mov (word [rdi + DR * 2], r8w);
                    xor_ (r9w, r9w);
                    test (r8w, 0x8000);
                    jnz (".n");
                    cmp (r8w, r9w);
                    jz (".z");
                    jnz (".p");
                    L(".n");
                    add (r9w, 4);
                    jmp (".out");
                    L(".z");
                    add (r9w, 2);
                    jmp (".out");
                    L(".p");
                    add (r9w, 1);
                    jmp (".out");
                    L(".out");
                    mov (word [rdi + 20], r9w);
                });

                DEF_INST (LEA,*insn_ptr, 0xF000, 0xE000, {
                     // REG (DR) = (REG (R_PC) + imm9) & 0xFFFF;
                    mov (r8, imm9);
                    add (r8w, dx);
                    mov (word [rdi + DR * 2], r8w);
                    xor_ (r9w, r9w);
                    test (r8w, 0x8000);
                    jnz (".n");
                    cmp (r8w, r9w);
                    jz (".z");
                    jnz (".p");
                    L(".n");
                    add (r9w, 4);
                    jmp (".out");
                    L(".z");
                    add (r9w, 2);
                    jmp (".out");
                    L(".p");
                    add (r9w, 1);
                    jmp (".out");
                    L(".out");
                    mov (word [rdi + 20], r9w);
                });

                DEF_INST (NOT, *insn_ptr,0xF03F, 0x903F, {
                    mov (r8w, word [rdi + SR1 * 2]);
                    not_ (r8w);
                    mov (word [rdi + DR * 2], r8w);
                    xor_ (r9w, r9w);
                    test (r8w, 0x8000);
                    jnz (".n");
                    cmp (r8w, r9w);
                    jz (".z");
                    jnz (".p");
                    L(".n");
                    add (r9w, 4);
                    jmp (".out");
                    L(".z");
                    add (r9w, 2);
                    jmp (".out");
                    L(".p");
                    add (r9w, 1);
                    jmp (".out");
                    L(".out");
                    mov (word [rdi + 20], r9w);
                }); 

                DEF_INST (RTI,*insn_ptr, 0xFFFF, 0x8000, {
                    // push (rdi);
                    // push (rsi);
                    // push (rdx);
                    // push (rcx);
                    // // read_memory : (reg ms pc insn)->memrd
                    // mov (rdi, INSN_RTI);
                    // mov (rax, (size_t)helper_dbg);
                    // call (rax);
                    // pop (rcx);
                    // pop (rdx);
                    // pop (rsi);
                    // pop (rdi);
                });
                
                DEF_INST (ST,*insn_ptr, 0xF000, 0x3000, {
                     //read_memory ((REG (R_PC) + imm9) & 0xFFFF, &(REG (DR)));
                    push (rdi);
                    push (rsi);
                    push (rdx);
                    push (rcx);
                    // read_memory : (reg ms pc insn)->memrd
                    mov (rsi, (uint64_t)ms);
                    call (helper_st);
                    pop (rcx);
                    pop (rdx);
                    pop (rsi);
                    pop (rdi);
                });

                DEF_INST (STI, *insn_ptr,0xF000, 0xB000, {
                    push (rdi);
                    push (rsi);
                    push (rdx);
                    push (rcx);
                    // read_memory : (reg ms pc insn)->memrd
                    mov (rsi, (uint64_t)ms);
                    call (helper_sti);
                    pop (rcx);
                    pop (rdx);
                    pop (rsi);
                    pop (rdi);
                });

                DEF_INST (STR,*insn_ptr, 0xF000, 0x7000, {
                    push (rdi);
                    push (rsi);
                    push (rdx);
                    push (rcx);
                    // read_memory : (reg ms pc insn)->memrd
                    mov (rsi, (uint64_t)ms);
                    call (helper_str);
                    pop (rcx);
                    pop (rdx);
                    pop (rsi);
                    pop (rdi);
                });

                DEF_INST (TRAP,*insn_ptr, 0xFF00, 0xF000, {
                    mov (word [rdi + 14], dx);
                    mov (r8w, word [rsi + vec8 * 2]);
                    mov (rdx, r8);
                    mov (rax, rdx);
                });
                executed:
                outLocalLabel ();
            }
            // sync pc
            mov (word [rdi + 18], rdx);
            ret ();
        }
};

void *
jit_exec_malloc (size_t size) {
    //using namespace Xbyak;
    #define PGSZ 4096
    size_t size_aligned;
    if (size % PGSZ == 0) {
        size_aligned = size;
    } else {
        size_aligned = (size / PGSZ + 1) * PGSZ;
    }
    uint8_t * mem = (uint8_t*)malloc (size_aligned + 16);
    uint8_t *mem_aligned = Xbyak::CodeArray::getAlignedAddress (mem);
    return  mem_aligned;
}

void *
jit_translate (void * blk, void *exec_ptr, size_t size) {
    //using namespace Xbyak;
    LC3Generator gen ((BasicBlock*)blk, exec_ptr, size);
    if (!Xbyak::CodeArray::protect (exec_ptr, size, Xbyak::CodeArray::PROTECT_RWE)) {
        fprintf(stderr, "can't protect\n");
        exit (-1);
    }
    void * native_code = gen.getCode <void *> ();
    //printf ("genearted.\n");
    return native_code; 
}